feat(rl): add shared-prefix packing and hybrid model execution - #7914

Remove the numerical-discrepancy discussion while retaining the native distributed and kernel replay qualification requirements.

Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

Retain the production attention arithmetic and public diagnostic entry points.
Consolidate tensor planning helpers, remove inactive MTP helpers, and defer
three unused default-off Triton attention variants with explicit import guards.

Add registered native attention replay coverage with independent references,
stream contention, and real NCCL CP1/2/4 tests. Keep the unchanged Mamba
implementation and document outstanding full-model and determinism gates.

Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

Remove TransformerConfig.sequence_relative_kernels, ModelParallelConfig.
deterministic_tp_reduce_scatter, the ordered TP reduce-scatter, the
sequence-relative attention adapter and Mamba scan, their TE/Mamba/attention
hooks, manifest entry and dispatch test. Neither NeMo-RL nor the shared-prefix
path uses them; the reverse patch is kept for a separate batch-invariant PR.
Removing the fields also restores the HybridModel GOLDEN_CONFIG field set and
drops the two auto-generated CLI flags.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

Attention read _shared_prefix_forest via getattr on an undeclared attribute
injected by the shared-prefix stack. Declare a class-level None default so the
writer's per-call set/del falls back to it, read it as a plain attribute, and
pin the protocol in test_attention.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

Run tensor_parallel.checkpoint and CheckpointWithoutOutput on CUDA with a
ContextVar and a tensor observer set only during forward; assert recompute on
the autograd thread sees the forward scope, observes once, and leaks nothing.
Fails at base (recompute sees the default scope).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

TopKRouter.forward always passed token_multiplicities= to self.routing, so any
replacement routing callable with the upstream signature raised TypeError on
the default path (40 upstream test failures). Forward it only when not None,
matching MoELayer.route, and let InferenceTopKRouter's training fallback
accept and forward it.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

_apply_expert_bias kept the legacy counting code next to the new
multiplicity-aware helper. Fold both into _expert_bias_token_counts, which
adds in place into the expert-bias buffer and keeps the upstream kernels
(bool sum, index_add_/scatter_add_ by determinism) when no multiplicities are
given. Multiplicities are cast to the int64 buffer dtype so logical counts
stay exact beyond 2^24.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

The shared-prefix fixed-row router GEMM padded and copied every 1024-row
block. Pad only the tail once and run each block on a view of the padded
input (bitwise identical; ~1.9x faster at 16k and 64k rows on GB200). Outside
the shared-prefix scope, forward calls the single GEMM directly.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

…ests

Register mamba_ragged.py/mamba_ragged_scan.py (triton), mamba_sequence_packing.py and
shared_prefix.py (dispatch) in the kernel manifest, and add bit-exact replay tests for the
ragged gather, the forest scan (with a MAMBA_DETERMINISTIC=0 negative control) and the
ragged/packed scans through a real MambaMixer.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

causal_conv1d's channel-first backward returns wrong dx/dw/db when
L % 1024 is 1..7 (bf16/fp16) or L % 512 is 1..3 (fp32). The state_fork,
replay_prefix and CP1 single-star helpers transposed xbc to channel-first
storage; keep [batch, length, channels] and pass a channel-last view, as
MambaMixer and the ragged path do. fp32 grads at conv lengths 513/1025
vs dense: 3.5e-3 -> 1.5e-7 (IEEE GEMMs).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

TP1/CP1 single stars bypassed NRL_SP_MAMBA_IMPL through a dedicated
padded path, and the default state_fork pads every branch to the longest
sibling (G16 skewed: 49 ms vs 8.4 ms ragged, 12.4 vs 1.25 GiB). Route
every Mamba layer through the selector, make ragged_state_fork the
default, delete the now-unreachable single-star helpers, and validate
NRL_SP_MAMBA_IMPL plus ragged's preconditions (power-of-two chunk_size,
default SSM state dtype) before any layer runs.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

The ragged gather backward specialized on the pack length and the forest
state kernels on every stride, so each new length or chunk count JIT
compiled new kernels (gather_bwd 1 -> 18, state_fwd 2 -> 20 variants;
58-108 ms per new length, 350-580 ms per new chunk count). Store the
contributor map row-major so the gather needs no length argument, and
pass strides at runtime. Offsets switch to int64 through a constexpr
index dtype only when a tensor spans 2^31 elements, which removes the
illegal memory accesses above 2^31 and keeps kernel times unchanged.

The ragged gather replay test's sanity check now counts the copies of
each input row (dim 1 of the row-major map), so it still requires a
row that is copied into a sibling branch.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

ragged_state_fork, now the default, imports private mamba_ssm SSD
kernels lazily on the first Mamba forward. Import them while validating
the stack so a mamba_ssm without that surface fails before any layer
runs, with a message naming the state_fork fallback.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

_validate_hybrid_stack now takes the global physical star length instead of
the hidden states, so a caller can run it before the embedding or any other
compute. The forward entry checks the hidden-state dtype/shape and derives the
physical length once, removing its second derivation.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

Reject contiguous linear CP layouts at CP>1 (the shared Mamba path relies on
the zigzag undo), MoE rank-capacity token dropping and routing replay. Accept
aux-loss load-balancing types whose coefficients are zero (the MCore default)
and reject only routing-changing types (sinkhorn, quantile_balancing) or
active auxiliary losses. Add CPU guard tests.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

forward_hybrid_stack_shared_prefix no longer infers whether per-branch padding
counts toward expert-bias statistics from the dispatcher type. The new
exclude_sequence_padding_from_expert_bias keyword defaults to the MCore dense
convention (padding counted, as with padding_mask=None); callers that mask
packed-sequence padding in their dense path pass True.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

…iers

The module docstring and capability comments called topologies and the MTP
dense-heads feature "validated" and the prefix replay a fallback. Reword them
as implemented code paths and contract identifiers; token strings unchanged.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

…edicate

The helper was equivalent to the base predicate on every reachable state
(mtp_num_layers=0 with an MTP pattern already raises in __init__), and its
docstring described that unreachable case.

Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

The grouped branch of process_mtp_loss split the rolled loss_mask instead of
layer_loss_mask, so with mtp_input_mask the per-group denominator ignored the
cumulative input mask and broke the "as if separate forwards" contract.

Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Training-mode shared-prefix MTP rebuilds dense branches from input_ids and
loss_mask only, so reject decoder_input and mtp_input_mask up front instead
of diverging from the dense path, and require calculate_per_token_loss for
every forest (each root is its own loss group) before the backbone runs.
Gate these checks on compute_mtp_loss so callers can skip MTP explicitly,
drop the constant-false MTP capability check, document the eval-mode MTP
skip under compute_mtp_loss, and move the shared_prefix_layout contract into
the forward Args section.

Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…scope

The backbone's router_gating_token_blocks scope ended before the MTP block,
so MoE layers in the MTP predictor routed with pack-size-dependent GEMMs.

Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…and TP/CP MTP

Mamba matrix: a prompt shorter than the convolution halo and D_has_hdim=True
for every backend. Distributed HybridModel parity: also run ragged_state_fork
(the documented NeMo-RL backend) at TP2/SP/CP2. MTP: shared-vs-dense parity at
TP2/SP/CP2 (lengths with unit CP count ratios, so normalization is comparable),
which exercises the CP/TP branch repacking of _forward_shared_prefix_mtp; the
gradient metric can now be restricted to the MTP heads across TP shards.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

The model-level MTP tests used lengths whose MTP count ratios are all 1, where every
grouping gives the same weights. The new test records the per-token normalized MTP loss
and cross entropy of a HybridModel forest forward whose ratios differ from 1 (1- and
2-token prompts, CP2 zigzag), and compares norm / CE with the analytic
factor / D * original / rolled weight of each layout loss group, per CP rank and depth.
It covers the default per-root groups and explicit root counts (2, 1) at TP1/CP1 and
TP2/SP/CP2. It asserts that the whole-pack, per-root and per-row groupings differ.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

…ne layout

Runs dense rows through HybridModel with shared_prefix_layout omitted and with
shared_prefix_layout=None, with and without MTP heads, and requires bitwise equal
logits, parameter gradients and router token counts. Deterministic kernels are enabled
for the test, because the Mamba backward is otherwise not bit-reproducible
between two identical calls.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

HybridModel.forward now runs its shared-prefix entry checks before the
fine-grained-offloading and paged-stash preprocessing, the embedding and
the RoPE gather. It derives the global physical length once (from
input_ids, or from a caller-supplied decoder_input) and passes it to
_validate_hybrid_stack, so a rejected configuration fails before any
state is mutated or any partial forward runs.

The model-level copies of checks the stack validator already performs
are removed: the TP/SP pairing, the TP config/process-group match, the
_validate_shared_prefix_physical_length call and the decoder_input shape
re-check in the RoPE branch. The model-state checks (cp_batch, inference,
masks, labels, packed_seq_params, PP1, MTP inputs, position embedding,
MLA, parallel output) stay.

In shared_prefix.py, _validate_mamba_fork no longer repeats the TP/SP
pairing that the stack validator checks, and the Mamba layer forwards
and the packed_fused oracle no longer re-run validation on every layer
call (and again under recompute). The stack validator still calls
_validate_mamba_fork for every Mamba layer before any layer executes.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

…onvention

HybridModel.forward gains the keyword-only argument
shared_prefix_exclude_sequence_padding_from_expert_bias: bool = False and
passes it to forward_hybrid_stack_shared_prefix as
exclude_sequence_padding_from_expert_bias. False, the MCore dense
convention, counts per-branch padding rows in the expert-bias token counts,
as a dense forward with padding_mask=None does; True excludes them, as a
dense forward with the packed-sequence padding mask does. Passing True
without shared_prefix_layout raises ValueError. The model and the stack
function share this default, and neither infers the convention from the
dispatcher type.

Tests: a TP1/CP1 MoE HybridModel with per-branch padding and replayed
routing reproduces the dense counts with padding_mask=None (flag unset or
False) and with the packed padding mask (flag True), and passing the flag
without a layout raises.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

MoELayer now declares _shared_prefix_token_multiplicities = None at class
level, with a comment naming its writer, and reads it with plain
attribute access, as Attention already does for _shared_prefix_forest.
forward_hybrid_stack_shared_prefix reads both slots the same way, so a
refactor that renames either one fails loudly instead of silently
dropping the tree mask or the multiplicities. The writer's per-call set
and del keep working: del falls back to the class default.

The router_gating_token_blocks docstring no longer says the scope must
be entered inside every checkpoint callable: Megatron's
tensor_parallel.checkpoint and CheckpointWithoutOutput restore the
forward context during recompute; only te_checkpoint and
torch.utils.checkpoint need the scope inside the callable.

The guard test's Mock(spec=MoELayer) now sets the declared slot to None,
as the class default does.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

…r effect

_validate_hybrid_stack compared window_size against (None, (-1, -1)), so
the full-attention window written as a YAML/CLI list [-1, -1] was
rejected as sliding-window attention. It now compares tuple(window_size).

The auxiliary-loss check rejected any nonzero moe_aux_loss_coeff, even
when the paired load-balancing type is "none", where TopKRouter adds no
loss. A Nemotron-H config with the default coefficient and type "none"
passed NeMo-RL's setup check but failed at the first shared-prefix
forward. Coefficients are now checked per (type, coefficient) pair and
rejected only for the auxiliary-loss types, matching
TopKRouter.get_aux_loss_coeff.

When the stack has attention layers, validation also checks the
flash-attn >= 2.7.0 requirement of the fused attention backward (a
missing flash-attn fails the same way). The attention entry point keeps
its own check; this one fails before the first Mamba layer runs.

Tests: list and tuple full windows, aux coefficients paired with "none",
the flash-attn gate, and a rejection matrix over the remaining stack
guards (hash routing, quantization, wide residual, mHC, softcapping,
TP/SP/CP/PP topology, recompute, offloading, CUDA graphs, fp8/fp4,
dropout, sliding window, softmax type, forced routing, z-loss, jitter,
MLP chunking, expert bias without logical lengths, layer types,
selective core-attention recompute, QK clipping).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

_forward_mamba_layer_shared_prefix_cp, which the module docstring names
as the backend selector, now states that ragged_state_fork is the
default for every topology and root count, the numerical, determinism
and performance evidence behind it (with the test files), its private
mamba_ssm dependency and the state_fork workaround, and what each opt-in
value is for. The packed_fused oracle and the prefix-replay helper are no
longer called fallbacks: one is a diagnostic isolation oracle, the other
a parity baseline.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

…_MAMBA_IMPL

The default-backend change removed the TP1/CP1 single-star Mamba helpers,
so the numerics matrix's "star_cp1" backend now always skipped. It is
removed, together with its star-only special cases, and MAMBA_BACKENDS
lists the four differentiable NRL_SP_MAMBA_IMPL values. The backend
helper no longer turns an unknown backend name into a skip: it fails, so
a renamed backend cannot silently drop out of the matrix.

The distributed model-parity case that pinned NRL_SP_MAMBA_IMPL to
ragged_state_fork duplicated the unset (default) case; it now runs the
opt-in state_fork backend at TP2/SP/CP2 instead, and the default cases
clear the variable explicitly.

test_shared_prefix_mamba_backends.py drops its causal_conv1d length test:
the numerics file covers the same channel-first defect for every backend
at more lengths (513..515, 1025..1031) with IEEE FP32 and BF16 checks,
and it was verified to fail without the channel-last fix. The backends file keeps the
selector, validation, dispatch and ragged-kernel tests; its fake mixer no
longer carries TP/SP fields that validation stopped reading.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

…resh manifest

test_shared_prefix_attention.py gains a module-level replay of the
shared-prefix branch of SelfAttention.forward: a real TE SelfAttention
with the forest attached through Attention._shared_prefix_forest, RoPE
positions and CP shards built as HybridModel builds them, at CP1 and CP2,
replaying the output, the input gradient and every parameter gradient
under torch deterministic algorithms. The file is now listed for the
attention_kernels_and_dispatch entry, whose notes name the branch.

Manifest updates:
- shared_prefix_composed_attention: the notes described the removed
  NRL_SP_DETERMINISTIC_BACKWARD switch and Triton import guards. They now
  state that torch.use_deterministic_algorithms(True) alone gives exact
  replay, the flash-attn >= 2.7.0 keyword backward and the 2x dense-branch
  error bound.
- ssm_mamba_ragged_forest, ssm_mamba_sequence_packing and
  shared_prefix_hybrid_dispatch also cite the Mamba backends, Mamba
  numerics and HybridModel parity tests, and say what each checks.
- shared_prefix_hybrid_dispatch: the notes now describe the
  NRL_SP_MAMBA_IMPL selector with ragged_state_fork as the default and
  state_fork/replay_prefix as opt-in helpers without a bit-exact backward
  replay.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

The PR changes moe_utils.py, router.py and moe_layer.py (fixed router
row blocks, logical token multiplicities for expert-bias counts, the
MoELayer multiplicity slot), but test_moe_kernels.py never exercised
those paths, so the CI determinism gate flagged all three entries.

New replays, each with three runs under stream contention:
- router_gating_linear inside router_gating_token_blocks(1024) on 8492
  rows (uneven tail), FP32/FP64 router dtype, with and without bias;
- TopKRouter with expert bias and token_multiplicities, comparing the
  probabilities, the routing map, all gradients and the int64
  expert-bias counts of the call (which must equal topk * sum of the
  multiplicities);
- MoELayer (allgather and alltoall dispatchers) under the shared-prefix
  scope forward_hybrid_stack_shared_prefix applies: multiplicities on
  the declared slot and the router row-block scope, with deterministic
  algorithms, again including the expert-bias counts.

tools/check_kernel_determinism_coverage.py --base-ref base now reports
no issues for the branch.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

validate_shared_prefix_execution_units now also requires every shared
unit to execute exactly its layout: the layout's rows in layout order
and the layout's physical total length. Planner and alignment outputs
already satisfy this; a hand-built or stale unit now fails at
validation instead of materializing rows in the wrong order.

The pack_dense_fallbacks error no longer calls the mode "DP1": NeMo-RL
allows pack_groups at DP>1 with align_data_parallel. The repack comment
is reworded the same way.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

plan_fixed_execution_slots rejects a sequence_length_pad_multiple that
is not a positive integer with a clear ValueError (None used to fail
with a TypeError from the comparison).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

make_repeated_group_ids had no production caller in Megatron-LM or
NeMo-RL. The test that used it now spells out its group IDs, the two
tests of the helper itself are gone, and the NeMo-RL API contract no
longer lists it.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

megatron/training/training.py imported megatron.rl.rl_utils and
rl_profiling at module import, guarded only by ImportError. The PR's
narrower wheel makes that import fail cleanly for wheel installs, but
NeMo-RL installs megatron-core in editable mode, where setuptools'
editable finder maps the whole megatron/rl directory. Every NeMo-RL
Megatron worker therefore still set has_rl_utils=True and imported the
RL agent, inference and server stack (fastapi, uvicorn, openai) without
running an RL step.

The module now binds placeholder names and _import_rl_utils(args) loads
megatron.rl on first use when perform_rl_step, rl_use_sequence_packing or
rl_profile is set; pretrain() and train() call it before any RL code
path. A missing package still leaves has_rl_utils False, and the
existing "RL cannot run without the megatron.rl package" assertion is
unchanged. Missing attributes on mock args count as disabled.

Test: a subprocess imports megatron.training.training, checks that
megatron.rl.rl_utils is not loaded and that a non-RL args object keeps
it unloaded, then that an RL args object loads it and binds the
profiler names.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

The megatron-core wheel now ships megatron/rl's shared-prefix packing
modules, but no gate covered them: pre-commit and tools/autoformat.sh
lint only megatron/core and tests, and the install test checks only
megatron.core and megatron.training imports.

- .pre-commit-config.yaml: black, isort and pylint also run on
  megatron/rl/shared_prefix_*.py and megatron/rl/tree_layout.py.
- tools/autoformat.sh: the same files join CHANGED_FILES.
- install-test.yml: the pip job imports megatron.rl.shared_prefix_execution
  and megatron.rl.shared_prefix_tensors from the installed package.

To pass the new pylint gate (docstring checks), the forest-layout
properties, ExecutionPlan properties, AlignedUnit and four PackedTreeLayout
members gain one-line docstrings. black, isort and ruff were already
clean on these files; pylint is now 10.00/10 on them.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

The ssm extra now notes that the default shared-prefix Mamba backend
(ragged_state_fork) calls private mamba_ssm SSD kernels, validated with
mamba-ssm 2.3.2.post1, that shared-prefix validation fails early when
they cannot be imported, and that NRL_SP_MAMBA_IMPL=state_fork uses only
the public API. The version is not pinned here: mamba-ssm already
resolves to a fixed git revision through [tool.uv.sources], and a
specifier change would need a uv.lock refresh.

The py-modules comment no longer says that shipping the RL runtime turns
on megatron.training's RL import: training.py now imports megatron.rl
only for RL runs, for wheel and editable installs alike.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

… feature

test_shared_prefix_port_contracts.py mixed four mock tests that pinned
the keyword arguments of MoELayer.route, MoELayer.forward,
TopKRouter.forward and TopKRouter.routing with unrelated feature tests.
The mock tests are deleted: the behavior they stood in for is now tested
directly (router forwarding of multiplicities only when set in
test_router_shared_prefix.py; logical expert counts through a real
MoELayer, with recompute, in test_shared_prefix_moe_mtp.py; HybridModel
counts in test_shared_prefix_model_parity.py; bit-exact MoE replays in
test_moe_kernels.py).

The remaining tests move next to the code they test:
- invalid expert-count metadata and the router GEMM backward with frozen
  inputs -> tests/unit_tests/transformer/moe/test_router_shared_prefix.py;
- Mamba branch pack/merge -> tests/unit_tests/ssm/test_mamba_branch_layout.py;
- layout multiplicities and positions ->
  tests/unit_tests/models/hybrid/test_shared_prefix_layout.py;
- recompute context restore without duplicate observations ->
  tests/unit_tests/tensor_parallel/test_random.py.

test_logical_expert_counts_match_expanded_rows is dropped:
test_expert_bias_counts_accumulate_in_place in test_router_shared_prefix.py
checks the same dense and boolean routing maps, padding and multiplicities
against a closed-form reference.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

TestSharedPrefixMTPDistributedParity gains a reproducibility check for
the model the MTP fixes target: a two-root forest through HybridModel
with MoE backbone layers and two MoE MTP predictor depths at TP2/SP/CP2,
with natural (not replayed) routing and deterministic kernels. Two
shared forwards/backwards must give identical completion logits, every
parameter gradient (MTP heads included) and the expert-bias token counts
of all four routers. The existing reproducibility test covers only an
attention/MLP model at TP1/CP1.

Passes on GB200 (4 ranks).

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

…assed

hybrid_model.py imported megatron.core.models.hybrid.shared_prefix at
module scope. That module imports einops unconditionally and loads the
Mamba layer and mixer modules, so importing HybridModel (and with it
llava_model and megatron.core.models.mamba) failed without einops even
for attention/MoE-only models that never pass shared_prefix_layout.
einops is not a core dependency, and mamba_mixer.py guards its own
import of it.

Import only the layout types at module scope, from shared_prefix_layout,
which needs only torch, and import _validate_hybrid_stack and
forward_hybrid_stack_shared_prefix inside the shared-prefix branches.
Tests that replace the stack forward now patch it on the shared_prefix
module. A new test imports hybrid_model in a fresh interpreter with
einops (and Transformer Engine, which needs it) hidden, and checks that
the adapter module is not loaded.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

…ward

HybridModel.forward validates the stack before its embedding, and
forward_hybrid_stack_shared_prefix then validated it again: every layer
was revisited, the ragged Mamba kernel import retried per Mamba layer and
the flash-attn version checked twice. The two calls also derived the
global physical length separately, from input_ids and from hidden_states.

forward_hybrid_stack_shared_prefix takes a keyword-only
validated_physical_len. When a caller passes the length it already
validated the stack and layout for, the stack function skips the
validation and only checks that hidden_states implies the same length.
HybridModel passes its length. Direct callers of the stack function are
unchanged and still validated.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

…tests by feature

Two tests checked that one grouped process_mtp_loss call equals one call
per loss group, with and without mtp_input_mask, each with its own
fixture. Keep test_grouped_mtp_loss_matches_independent_forwards in
test_shared_prefix_moe_mtp.py, whose pack also puts loss-mask starts at
segment boundaries, and delete the copy and its helper from the other
file.

Rename test_shared_prefix_mtp_fixes.py, which now holds the fixed-order
prompt-row gather tests and the model-level MTP input and router-scope
tests, to test_shared_prefix_mtp_branches.py.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

…on_slots

plan_shared_prefix_bins and the oversized-slot split already spread an
exact-prompt group larger than the branch limit evenly (17 -> 9+8). The
fixed-slot packer did not: its FFD filled each bin up to
max_rows_per_slot, so a group of 17 became a 16-row slot plus a
singleton. That singleton ran as an unshared dense forward both in the
driver's slot plan and in plan_shared_prefix_execution_units with
pack_groups and repack_groups, which replans each group with this
packer. Logprobs stayed exact; only the sharing was lost.

Cap the rows per bin at ceil(n / ceil(n / max_rows_per_slot)), so the
packer uses the same number of bins as before when token capacity does
not bind and fills them evenly. Tests cover 16, 17, 32 and 33 rows for
the slot planner directly and for the repack path.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

…refix docstrings

The module docstring cited two commits of an internal fork that readers
of this repository cannot reach, and the Mamba backend docstring quoted
GB200 speedups. Drop the commit IDs, and state why ragged_state_fork is
the default (it pads each branch only to its own chunk boundary) in
place of the measured ratios, which belong in the PR description.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

The developer notes still described controls and evidence that no longer
exist. Remove the sequence-relative/deterministic-reduce-scatter
paragraph (that feature moved to a separate change) and the paragraph on
the deleted NRL_SP_FUSED_* opt-ins, and rewrite the validation status
around the current test files instead of stale case counts.

Also document what changed in the code: model execution lives in
megatron.core.models.hybrid and is imported only with a layout; the
attention plan and its deterministic backward; the NRL_SP_MAMBA_IMPL
values with ragged_state_fork as the default, its requirements and the
Mamba kernel replay; the expert-bias padding keyword (padding counted by
default) and exact int64 logical counts; the training-only MTP contract,
its fixed-order prompt gradient and per-group normalization; the full
list of rejected features, including per-rank capacity dropping, routing
replay and the contiguous Mamba CP layout, and the acceptance of
zero-coefficient auxiliary-loss balancing; and capability tokens as
contract identifiers. Add the shared-prefix forest attention to the
determinism op catalog.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

Drop references to APIs that no longer exist: the
materialize_attention_mask argument, the stable inverse row order and
the repeated rollout-group identities. The reference attention masks now
live in the test oracles. Describe rank_order_permutation as the forward
permutation that restores source order.

Document MAX_SHARED_PREFIX_BRANCHES and the even split of larger
exact-prompt groups, the equal-group-size rule and the cases that run a
single row as a forest root, the determinism contract for dense_packer,
and the TP/CP padding quantum. State the narrowed wheel scope (the
packing modules, generation_api and the lazy __init__), noting that an
editable install still exposes all of megatron/rl, and add
test_tree_layout.py to the test command.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

…notes

Add the shared-prefix developer notes to the Developer Guide toctree so
the docs build reaches them, and link the packing API notes, which live
outside the Sphinx source tree, by their repository URL.

Describe the validation method in terms of what the tests check: FP32
exactness against local PyTorch layers with TF32 disabled, and why
Transformer Engine FP32 references need NVIDIA_TF32_OVERRIDE=0. Explain
the BF16 shared-versus-dense gradient difference as rounding amplified
by MoE top-k flips, and state that a repeated dense run measures only
nondeterminism, so the null is a dense run with a different packing or
the distance from a routing-replayed high-precision reference. State
that no whole-model comparison at a large TP/CP topology is part of the
suite.

Also note the channel-last causal_conv1d input on every shared Mamba
path, restate the MTP prompt-gather memory bound plainly, and say that
NeMo RL imports megatron.rl directly instead of re-exporting it.

Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…and contiguous CP

The three rejections now state the reason in the error: the expert rank
capacity budget counts physical star rows and a dropped prompt row is lost
for every completion that shares it; routing-replay targets are indexed by
dense rows; and the shared-prefix forward relies on the zigzag CP layout.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>

… does not state it

Unset, per-branch padding rows are excluded from expert-bias counts with the flex
dispatcher's HybridEP backend and counted otherwise; an explicit True or False wins.
Stating the convention without shared_prefix_layout still raises ValueError.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Signed-off-by: Jorge Albericio <jalbericiola@nvidia.com>